{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial: 基于CPM-Bee进行解码层微调（decoder tuning）\n",
    "本教程基于情感分类数据集SST2对CPM-Bee进行解码层微调（decoder tuning）。解码层微调（decoder tuning）是在不训练模型的情况下，通过加入输出端的解码器网络，使用少样本训练解码器网络来提升模型的理解能力。我们将16 shot的微调结果与原始模型zero shot进行对比。\n",
    "\n",
    "This tutorial is based on the sentiment classification data set SST2 for CPM-Bee decoder tuning. decoder tuning is to improve the understanding ability of the model by joining the decoder network at the output end and training the decoder network with few samples without training the model. We compared the fine-tuning results of 16 shot to the original model zero shot."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. 数据格式处理 (Process dataset)\n",
    "训练之前，我们需要定义并处理我们的数据输入格式，我们构造一个数据集的处理类，将数据处理为特定格式。\n",
    "\n",
    "Before training, we need to define and process our data input format. We construct a processing class for the data set to process the data into a specific format."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在本教程中，我们使用的情感分类的输入格式如下（也可以自行定义其他格式）：\n",
    "\n",
    "In this tutorial, we use the following input format for emotion classification (you can also define other formats) :\n",
    "```\n",
    "input: text\n",
    "question: \"What is the sentiment of this sentence?\"\n",
    "<ans>: bad/great \n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "class SST2Processor():\n",
    "    \"\"\"\n",
    "    `SST-2 <https://nlp.stanford.edu/sentiment/index.html>`_ dataset is a dataset for sentiment analysis. It is a modified version containing only binary labels (negative or somewhat negative vs somewhat positive or positive with neutral sentences discarded) on top of the original 5-labeled dataset released first in `Recursive Deep Models for Semantic Compositionality Over a Sentiment Treebank <https://aclanthology.org/D13-1170.pdf>`_\n",
    "\n",
    "    We use the data released in `Making Pre-trained Language Models Better Few-shot Learners (Gao et al. 2020) <https://arxiv.org/pdf/2012.15723.pdf>`_\n",
    "\n",
    "    \"\"\"\n",
    "    dataset_project = {\"train\": \"train\",\n",
    "                        \"dev\": \"dev\",\n",
    "                        \"test\": \"test\"\n",
    "                        }\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.labels = ['0', '1']\n",
    "        self.verbalizer = {\"0\": \"bad\", \"1\": \"great\"}\n",
    "        self.label_word = [\"bad\",\"great\"]\n",
    "    \n",
    "    def get_examples(self, data_dir, split,shot =-1):\n",
    "        self.counts = {\"bad\":0,\"great\":0}\n",
    "        path = os.path.join(data_dir, f\"{self.dataset_project[split]}.tsv\")\n",
    "        examples = []\n",
    "        with open(path, encoding='utf-8')as f:\n",
    "            lines = f.readlines()\n",
    "            for idx, line in enumerate(lines[1:]):\n",
    "                linelist = line.strip().split('\\t')\n",
    "                text_a = linelist[0]\n",
    "                label = linelist[1]\n",
    "                guid = \"%s-%s\" % (split, idx)\n",
    "                example = {}\n",
    "                example[\"input\"] =text_a\n",
    "                example[\"question\"] = \"What is the sentiment of this sentence?\"\n",
    "                example[\"<ans>\"] = self.verbalizer[label]\n",
    "                if shot==-1:\n",
    "                    examples.append(example)\n",
    "                else:\n",
    "                    if self.counts[example[\"<ans>\"]]>=shot:\n",
    "                        continue    \n",
    "                    else:\n",
    "                        examples.append(example)\n",
    "                        self.counts[example[\"<ans>\"]]+=1\n",
    "        return examples"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "添加不同的数据处理对象，以实现不同的数据格式。\n",
    "\n",
    "Add different data processing objects to implement different data formats."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROCESSORS = {\n",
    "    \"sst2\": SST2Processor\n",
    "}\n",
    "dataset_name = \"sst2\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "添加工作路径\n",
    "\n",
    "Add working path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import sys\n",
    "random.seed(123)\n",
    "sys.path.append(\"../src\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "预处理数据格式，并且按照预处理格式将处理好的数据存储为二进制文件。训练集和验证集选取SST2数据集的train和dev文件中的数据并构造16 shot数据，测试集选取的为SST2数据集的test文件。文件路径在 ./decoder_tuning_data/raw_data/\n",
    "\n",
    "Preprocesse the data format, and the processed data is stored as binary files. The training set and verification set select the data in the train and dev files of the SST2 data set and construct 16 shot data, and the test set selects the test file of the SST2 data set. File path in ./decoder_tuning_data/raw_data/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/liwentao/miniconda3/envs/py38/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Shuffle step 1/2: 100%|██████████| 32/32 [00:00<00:00, 28795.91it/s]\n",
      "Shuffle step 2/2: 100%|██████████| 1/1 [00:00<00:00, 3097.71it/s]\n",
      "Shuffle step 1/2: 100%|██████████| 32/32 [00:00<00:00, 72354.57it/s]\n",
      "Shuffle step 2/2: 100%|██████████| 1/1 [00:00<00:00, 3184.74it/s]\n",
      "Shuffle step 1/2: 100%|██████████| 1821/1821 [00:00<00:00, 327241.97it/s]\n",
      "Shuffle step 2/2: 100%|██████████| 1/1 [00:00<00:00, 304.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32\n",
      "{'input': 'a stirring , funny and finally transporting re-imagining of beauty and the beast and 1930s horror films', 'question': 'What is the sentiment of this sentence?', '<ans>': 'great'}\n",
      "{'0': 'bad', '1': 'great'}\n",
      "['bad', 'great']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "from cpm_live.dataset import build_dataset, shuffle_dataset\n",
    "\n",
    "import shutil\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "def build_bin_data(data,output_dir, dataset_path,dataset_name):\n",
    "    output_path = \"./decoder_tuning_data/bin_data/\"  +output_dir\n",
    "\n",
    "    with build_dataset(\"tmp\", \"data\") as dataset:\n",
    "        for item in data:\n",
    "            dataset.write(item) # reformat_data(item)\n",
    "    shuffle_dataset(\n",
    "        \"tmp\",\n",
    "        os.path.join(output_path, dataset_path),\n",
    "        progress_bar=True,\n",
    "        output_name=dataset_name\n",
    "    )\n",
    "    shutil.rmtree(\"tmp\")\n",
    "\n",
    "processor = PROCESSORS[dataset_name]()\n",
    "path = 'decoder_tuning_data/raw_data/' + dataset_name\n",
    "train_dataset = processor.get_examples(path,\"train\",shot=16)\n",
    "\n",
    "valid_dataset = processor.get_examples(path,\"dev\",shot=16)\n",
    "\n",
    "test_dataset = processor.get_examples(path,\"test\")\n",
    "if os.path.exists(\"./decoder_tuning_data/bin_data\"):\n",
    "    os.system(\"rm -rf ./decoder_tuning_data/bin_data/{}\".format(dataset_name))\n",
    "if os.path.exists(\"./tmp\"): \n",
    "    os.system(\"rm -rf ./tmp\")\n",
    "output_dir = dataset_name\n",
    "build_bin_data(train_dataset,output_dir, \"train_data\", \"example-data\")\n",
    "build_bin_data(valid_dataset,output_dir, \"valid_data\", \"example-data\")\n",
    "build_bin_data(test_dataset,output_dir, \"test_data\", \"example-data\")\n",
    "verbalizer = processor.verbalizer\n",
    "label_word = processor.label_word\n",
    "\n",
    "print(train_dataset[0])\n",
    "print(verbalizer)\n",
    "print(label_word)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. 训练（Training）\n",
    "我们自定义一个DecT_CPM类，方便实现在CPMBee上的decoder tuning，通过DecT_CPM中的run，run_zs函数分别实现16-shot下的decoder tuning和0-shot下的原本模型能力的测试。需要将预训练好的模型存储在`./ckpt/`文件夹下。具体而言，`./ckpt/`文件夹需要预先保存`./ckpt/config.json`、`./ckpt/pytorch_model.bin`和`./ckpt/vocab.txt`。\n",
    "\n",
    "We customize a DecT_CPM class to facilitate the realization of decoder tuning on CPMBee. Through the run and run_zs functions in DecT_CPM, decoder tuning under 4-shot and the original model ability test under 0-shot can be realized respectively. The pre-trained model needs to be stored in the './ckpt/ 'folder. Specifically, the './ckpt/ 'folder needs to be pre-stored'./ckpt/config.json ', './ckpt/pytorch_model.bin 'and'./ckpt/vocab.txt '."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "初始化一些超参数\n",
    "\n",
    "Initialize some hyperparameters and verbalizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 初始化一些超参数以及verbalizer\n",
    "lr = 4e-3\n",
    "proto_dim = 128\n",
    "model_logits_weight = 1\n",
    "max_epochs = 100"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "构造decoder_tuning的trainer，decoder tuning的主体模型是一个线性层，当数据量较大时，可以增加层数，以提升模型的学习能力。\n",
    "\n",
    "Construct the trainer of decoder_tuning. The main model of decoder tuning is a linear layer. When the amount of data is large, the number of layers can be increased to improve the learning ability of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, shutil\n",
    "import sys\n",
    "sys.path.append(\".\")\n",
    "import time\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from tqdm import tqdm\n",
    "import dill\n",
    "import warnings\n",
    "from typing import Optional\n",
    "from typing import Callable, Union, Dict, List\n",
    "try:\n",
    "    from typing import OrderedDict\n",
    "except ImportError:\n",
    "    from collections import OrderedDict\n",
    "from sklearn.metrics import accuracy_score\n",
    "from copy import deepcopy\n",
    "\n",
    "class DecTCPM(object):\n",
    "    r\"\"\"A runner for DecT\n",
    "    This class is specially implemented for classification.\n",
    "    Decoder Tuning: Efficient Language Understanding as Decoding : https://arxiv.org/pdf/2212.08408.pdf\n",
    "\n",
    "    Args:\n",
    "        model (:obj:`CPMBeeTorch`): One ``CPMBeeTorch`` object.\n",
    "        test_dataloader (:obj:`FinetuneDataset`): The dataloader to bachify and process the test data.\n",
    "        tokenizer (:obj:`CPMBeeTokenizer`): The tokenizer to process the word.\n",
    "        verbalizer (:obj:`Verbalizer`): The verbalizer to map the label to the word.\n",
    "        device (:obj:`torch.device`): The device to run the model.\n",
    "        calibrate_dataloader (:obj:`FinetuneDataset`, optional): The dataloader that has empty input, to modify the output logits. Defaults to None.\n",
    "        lr (:obj:`float`, optional): The learning rate. Defaults to 5e-3.\n",
    "        hidden_size (:obj:`int`, optional): The hidden size of the model. Defaults to 4096.\n",
    "        mid_dim (:obj:`int`, optional): The dimension of the proto vector. Defaults to 128.\n",
    "        epochs (:obj:`int`, optional): The number of epochs to train. Defaults to 5.\n",
    "        model_logits_weight (:obj:`float`, optional): The weight of the model logits. Defaults to 1.\n",
    "    \"\"\"\n",
    "    def __init__(self, \n",
    "                 model,\n",
    "                 test_dataloader,\n",
    "                 val_dataloader,\n",
    "                 tokenizer,\n",
    "                 verbalizer,\n",
    "                 device: Optional[Union[str, torch.device]] = \"cuda:0\",\n",
    "                 calibrate_dataloader: Optional[List] = None,\n",
    "                 lr: Optional[float] = 5e-3,\n",
    "                 hidden_size: Optional[int] = 4096,\n",
    "                 mid_dim: Optional[int] = 128,\n",
    "                 epochs: Optional[int] = 5,\n",
    "                 model_logits_weight: Optional[float] = 1,\n",
    "                 ):\n",
    "        self.model = model\n",
    "        self.val_dataloader = val_dataloader\n",
    "        self.test_dataloader = test_dataloader\n",
    "        self.calibrate_dataloader = calibrate_dataloader\n",
    "        self.loss_function = torch.nn.CrossEntropyLoss()\n",
    "        self.device = device\n",
    "        ids = []\n",
    "        for idx in range(len(verbalizer.items())):\n",
    "            ids.append(tokenizer.encode(verbalizer[str(idx)])[0][0])\n",
    "        self.label_list = list(verbalizer.values())\n",
    "        self.label_word_token_ids = []\n",
    "        for label_word in self.label_list:\n",
    "            self.label_word_token_ids.append(tokenizer.encode(label_word)[0][0])\n",
    "        self.ids = ids #nn.Parameter(torch.tensor(ids), requires_grad=False)\n",
    "        self.num_classes = len(self.ids)\n",
    "        self.lr = lr\n",
    "        self.mid_dim = mid_dim\n",
    "        self.epochs = epochs\n",
    "        self.model_logits_weight = model_logits_weight\n",
    "        self.hidden_dims = hidden_size\n",
    "        self.reset_parameter()\n",
    "    \n",
    "    # reset the parameters, useful when you want to test different random seeds\n",
    "    # self.head is a linear layer, if you want to use other models, you can modify it (useful when there are more data)\n",
    "    def reset_parameter(self):\n",
    "        self.head = nn.Linear(self.hidden_dims, self.mid_dim, bias=False)\n",
    "        w = torch.empty((self.num_classes, self.mid_dim)).to(self.device)\n",
    "        nn.init.xavier_uniform_(w)\n",
    "        self.proto = nn.Parameter(w, requires_grad=False)\n",
    "        r = torch.ones(self.num_classes)\n",
    "        self.proto_r = nn.Parameter(r, requires_grad=True)\n",
    "        self.optimizer = torch.optim.Adam([p for n, p in self.head.named_parameters()] + [self.proto_r], lr=self.lr)\n",
    "\n",
    "\n",
    "    # get the logits and hidden states of the model, specifically for cpmbee model, you can modify it for other models\n",
    "    def get_logits_and_hidden(self,data):\n",
    "        input_ids = torch.from_numpy(data[\"inputs\"]).cuda().to(torch.int32)\n",
    "        input_ids_sub = torch.from_numpy(data[\"inputs_sub\"]).cuda().to(torch.int32)\n",
    "        input_length = torch.from_numpy(data[\"length\"]).cuda().to(torch.int32)\n",
    "        input_context = torch.from_numpy(data[\"context\"]).cuda().bool()\n",
    "        input_sample_ids = torch.from_numpy(data[\"sample_ids\"]).cuda().to(torch.int32)\n",
    "        input_num_segments = torch.from_numpy(data[\"num_segments\"]).cuda().to(torch.int32)\n",
    "        input_segment_ids = torch.from_numpy(data[\"segment_ids\"]).cuda().to(torch.int32)\n",
    "        input_segment_rel_offset = (\n",
    "            torch.from_numpy(data[\"segment_rel_offset\"]).cuda().to(torch.int32)\n",
    "        )\n",
    "        input_segment_rel = torch.from_numpy(data[\"segment_rel\"]).cuda().to(torch.int32)\n",
    "        input_span = torch.from_numpy(data[\"spans\"]).cuda().to(torch.int32)\n",
    "        targets = torch.from_numpy(data[\"target\"]).cuda().to(torch.int32)\n",
    "        ext_table_ids = torch.from_numpy(data[\"ext_ids\"]).cuda().to(torch.int32)\n",
    "        ext_table_sub = torch.from_numpy(data[\"ext_sub\"]).cuda().to(torch.int32)\n",
    "        task_ids = torch.from_numpy(data[\"task_ids\"]).cuda().to(torch.int32)\n",
    "        task_names = data[\"task_names\"]\n",
    "        # to get the label from the targets\n",
    "        mask = torch.logical_or(targets ==self.ids[0], targets==self.ids[1])\n",
    "        labels = targets[mask]\n",
    "        final_label = []\n",
    "        for i in range(len(labels)):\n",
    "            final_label.append(self.ids.index(labels[i]))\n",
    "        with torch.no_grad():\n",
    "            logits, hidden_states = self.model(\n",
    "                    input_ids,\n",
    "                    input_ids_sub,\n",
    "                    input_length,\n",
    "                    input_context,\n",
    "                    input_sample_ids,\n",
    "                    input_num_segments,\n",
    "                    input_segment_ids,\n",
    "                    input_segment_rel_offset,\n",
    "                    input_segment_rel,\n",
    "                    input_span,\n",
    "                    ext_table_ids,\n",
    "                    ext_table_sub,\n",
    "                )\n",
    "        # mask the targets where value is -100 or 7, to get the index of the valid position\n",
    "        mask_matrix = deepcopy(targets)\n",
    "        mask_matrix[targets == -100] = 0\n",
    "        mask_matrix[targets == 7] = 0\n",
    "        index_mask = mask_matrix.nonzero(as_tuple=False)\n",
    "        # finally we get the logits and hidden states of the <ans> word position\n",
    "        filtered_logits = logits[index_mask[:, 0], index_mask[:, 1], :]\n",
    "        filtered_hiddens = hidden_states[index_mask[:, 0], index_mask[:, 1], :]\n",
    "        label_logits = filtered_logits[:,self.label_word_token_ids] # F.softmax(filtered_logits)[:,self.label_word_token_ids]\n",
    "        return label_logits, filtered_hiddens,final_label\n",
    "    \n",
    "    # test the model on the dev set, if zs is true, then test on the zero-shot setting, otherwise test on the decoder tuning setting\n",
    "    def test(self, dataloader,zs):\n",
    "        if zs:\n",
    "            preds = []\n",
    "            labels = []\n",
    "            for iteration, data in enumerate(dataloader):\n",
    "                if data is None:\n",
    "                    if last_data is None:\n",
    "                        raise RuntimeError(\n",
    "                            \"Dataset is too small, please use a smaller batch size or sequence length!\"\n",
    "                        )\n",
    "                    data = last_data  # use last data\n",
    "                    skip_this_batch = True\n",
    "                else:\n",
    "                    last_data = data\n",
    "                logits,_,label = self.get_logits_and_hidden(data)\n",
    "                preds.extend(torch.argmax(logits, dim=-1).cpu().tolist())\n",
    "                labels.extend(label)\n",
    "            res = sum([int(i==j) for i,j in zip(preds, labels)])/len(preds)\n",
    "            return res\n",
    "        else:\n",
    "            preds = []\n",
    "            labels = []\n",
    "            for iteration, data in enumerate(dataloader):\n",
    "                if data is None:\n",
    "                    if last_data is None:\n",
    "                        raise RuntimeError(\n",
    "                            \"Dataset is too small, please use a smaller batch size or sequence length!\"\n",
    "                        )\n",
    "                    data = last_data  # use last data\n",
    "                    skip_this_batch = True\n",
    "                else:\n",
    "                    last_data = data\n",
    "                logits,hidden_states,label = self.get_logits_and_hidden(data)\n",
    "                proto_logits = self.sim(self.head(hidden_states.float()), self.proto, self.proto_r, logits.float(), self.model_logits_weight).cpu()\n",
    "                preds.extend(torch.argmax(proto_logits, dim=-1).cpu().tolist())\n",
    "                labels.extend(label)\n",
    "            res = sum([int(i==j) for i,j in zip(preds, labels)])/len(preds)\n",
    "            return res\n",
    "\n",
    "    @staticmethod\n",
    "    def sim(x, y, r=0, model_logits=0, model_logits_weight=1):\n",
    "        x = torch.unsqueeze(x, -2)\n",
    "        x = F.normalize(x, dim=-1)\n",
    "        d = torch.norm((x - y), dim=-1)\n",
    "        dist = d - model_logits * model_logits_weight - r\n",
    "        return -dist\n",
    "    \n",
    "    # conduct the loss function in the decoder tuning\n",
    "    def loss_func(self, x, model_logits, labels):\n",
    "        sim_mat = torch.exp(self.sim(x, self.proto, self.proto_r, model_logits, self.model_logits_weight))\n",
    "        pos_score = torch.sum(sim_mat * F.one_hot(labels), -1)\n",
    "        loss = -torch.mean(torch.log(pos_score / sim_mat.sum(-1)))\n",
    "        return loss\n",
    "    \n",
    "    # run zero shot setting\n",
    "    def run_zs(self):\n",
    "        res = self.test(self.test_dataloader, zs = True)\n",
    "        print(\"zero shot acc:\",res)\n",
    "\n",
    "    # train the model with decoder tuning, you need to provide the training dataloader (type:FinetuneDataset)\n",
    "    def run(self, train_dataloader):\n",
    "        logits_list = []\n",
    "        hidden_states_list = []\n",
    "        labels = []\n",
    "        with torch.no_grad():\n",
    "            for iteration, data in enumerate(train_dataloader):\n",
    "                if data is None:\n",
    "                    if last_data is None:\n",
    "                        raise RuntimeError(\n",
    "                            \"Dataset is too small, please use a smaller batch size or sequence length!\"\n",
    "                        )\n",
    "                    data = last_data  # use last data\n",
    "                    skip_this_batch = True\n",
    "                else:\n",
    "                    last_data = data\n",
    "                train_logits, train_embeds,label = self.get_logits_and_hidden(data)\n",
    "                logits_list.append(train_logits)\n",
    "                hidden_states_list.append(train_embeds)\n",
    "                labels.extend(label)\n",
    "        train_logits = torch.cat(logits_list,dim=0)\n",
    "        train_embeds = torch.cat(hidden_states_list,dim=0)\n",
    "        embeds = [[] for _ in range(self.num_classes)]\n",
    "        train_labels = [[] for _ in range(self.num_classes)]\n",
    "        model_logits = [[] for _ in range(self.num_classes)]\n",
    "        total_num = 0\n",
    "        start_time = time.time()\n",
    "\n",
    "        for idx, label in enumerate(labels):\n",
    "            label = torch.tensor(label)\n",
    "            train_labels[label].append(label)\n",
    "            embeds[label].append(torch.tensor(train_embeds[idx]))\n",
    "            model_logits[label].append(torch.tensor(train_logits[idx]))\n",
    "        embeds = list(map(torch.stack, embeds))\n",
    "        labels = torch.cat(list(map(torch.stack, train_labels))).to(self.device)\n",
    "        model_logits = torch.cat(list(map(torch.stack, model_logits))).float()\n",
    "\n",
    "        self.head.to(self.device)\n",
    "        self.proto.to(self.device)\n",
    "        self.proto_r.to(self.device)\n",
    "        dist = list(map(lambda x: torch.norm(self.head(x.float()) - self.head(x.float().mean(0)), dim=-1).mean(), embeds))\n",
    "        self.proto_r.data = torch.stack(dist)\n",
    "        \n",
    "        loss = 0.\n",
    "        best_eval_res = 0.\n",
    "        \n",
    "        for epoch in range(self.epochs):\n",
    "            x = self.head(torch.cat(embeds).float())\n",
    "            self.optimizer.zero_grad()\n",
    "            loss = self.loss_func(x, model_logits, labels)\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            # use vaild dataset to evaluate the model, and test on best_eval_res\n",
    "            if epoch % 20 == 0 and epoch > 0 :\n",
    "                print(\"Total epoch: {}. DecT loss: {}\".format(epoch, loss))\n",
    "                eval_res = self.test(self.val_dataloader, zs = False)\n",
    "                print(\"val acc:\", eval_res)\n",
    "                if eval_res > best_eval_res:\n",
    "                    best_eval_res = eval_res\n",
    "                    test_res = self.test(self.test_dataloader, zs = False)\n",
    "                    print(\"test acc at best val:\",test_res)\n",
    "\n",
    "\n",
    "        end_time = time.time()\n",
    "        print(\"Total time: {}\".format(end_time - start_time))\n",
    "        res= self.test(self.test_dataloader, zs = False)\n",
    "        print(\"Final acc:\",res)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "加载模型的权重并部署\n",
    "\n",
    "Load the model's weights and deploy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cpm_live.tokenizers import CPMBeeTokenizer\n",
    "from cpm_live.training_tasks.bee import FinetuneDataset\n",
    "from cpm_live.models import CPMBeeConfig, CPMBeeTorch\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import bmtrain as bmt\n",
    "from copy import deepcopy\n",
    "config = CPMBeeConfig.from_json_file(\"./ckpt/config.json\")\n",
    "ckpt_path = \"./ckpt/pytorch_model.bin\"\n",
    "tokenizer = CPMBeeTokenizer()\n",
    "model = CPMBeeTorch(config=config)\n",
    "model.load_state_dict(torch.load(ckpt_path), strict=False)\n",
    "device = torch.device(\"cuda:0\")\n",
    "model.to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "构建dataloader\n",
    "build dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader = FinetuneDataset(\n",
    "        dataset_path = \"./decoder_tuning_data//bin_data/{}/train_data\".format(dataset_name),\n",
    "        batch_size=8,\n",
    "        max_length=512,\n",
    "        max_depth=8,\n",
    "        tokenizer=tokenizer,\n",
    "    )\n",
    "val_dataloader = FinetuneDataset(\n",
    "        dataset_path = \"./decoder_tuning_data/bin_data/{}/valid_data\".format(dataset_name),\n",
    "        batch_size=8,\n",
    "        max_length=512,\n",
    "        max_depth=8,\n",
    "        tokenizer=tokenizer,\n",
    "    )\n",
    "test_dataloader = FinetuneDataset(\n",
    "        dataset_path = \"./decoder_tuning_data/bin_data/{}/test_data\".format(dataset_name),\n",
    "        batch_size=8,\n",
    "        max_length=512,\n",
    "        max_depth=8,\n",
    "        tokenizer=tokenizer,\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "构建runner\n",
    "\n",
    "Build runner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "runner = DecTCPM(\n",
    "    model = model,\n",
    "    test_dataloader = test_dataloader,\n",
    "    val_dataloader=val_dataloader,\n",
    "    tokenizer = tokenizer,\n",
    "    verbalizer = verbalizer,\n",
    "    device = device,\n",
    "    calibrate_dataloader = None,\n",
    "    lr = lr,\n",
    "    mid_dim = proto_dim,\n",
    "    epochs = max_epochs,\n",
    "    model_logits_weight = model_logits_weight,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "开始训练，最后输出的是zero shot和decoder tuning的准确率\n",
    "\n",
    "Start training, the final output is zero shot and decoder tuning accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "zero shot acc: 0.8791872597473915\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2354165/3965538911.py:227: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  embeds[label].append(torch.tensor(train_embeds[idx]))\n",
      "/tmp/ipykernel_2354165/3965538911.py:228: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  model_logits[label].append(torch.tensor(train_logits[idx]))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total epoch: 20. DecT loss: 0.0790218934416771\n",
      "val acc: 0.9375\n",
      "test acc at best val: 0.9192751235584844\n",
      "Total epoch: 40. DecT loss: 0.05866580083966255\n",
      "val acc: 0.9375\n",
      "Total epoch: 60. DecT loss: 0.05533261597156525\n",
      "val acc: 0.9375\n",
      "Total epoch: 80. DecT loss: 0.0544731467962265\n",
      "val acc: 0.9375\n",
      "Total time: 36.721224308013916\n",
      "Final acc: 0.9203734211971444\n"
     ]
    }
   ],
   "source": [
    "runner.run_zs()\n",
    "runner.run(train_dataloader)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
